/* Copyright (C) 2013-2021 TU Dortmund
 * This file is part of AutomataLib, http://www.automatalib.net/.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package net.automatalib.commons.util.combinatorics;

import java.util.Arrays;

public class DisjointSetForestInt {

    private final int[] parent;
    private final int[] rank;

    public DisjointSetForestInt(int initSize) {
        this.parent = new int[initSize];
        Arrays.fill(this.parent, -1);
        this.rank = new int[initSize];
    }

    public int find(int a) {
        int p = parent[a];
        if (p == -1) {
            return a;
        }
        p = find(p);
        parent[a] = p;
        return p;
    }

    public int union(int a, int b) {
        return directUnion(find(a), find(b));
    }

    public int directUnion(int a, int b) {
        assert parent[a] == -1 && parent[b] == -1;

        if (a == b) {
            return a;
        }

        int ra = rank[a], rb = rank[b];
        if (ra < rb) {
            parent[a] = b;
            return b;
        }
        if (ra == rb) {
            rank[a]++;
        }
        parent[b] = a;
        return a;
    }

    public int size() {
        return parent.length;
    }

}
